package com.uh.rds.testing.utils;

import com.uh.rds.testing.data.FileTask;
import com.uh.rds.testing.data.TestDataGenerator;
import com.uh.rds.testing.base.DataValidator;
import com.uh.rds.testing.validator.DataValidatorRunner;
import com.uh.rds.testing.validator.types.*;
import org.junit.jupiter.api.Assertions;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.util.Map;

import static com.uh.rds.testing.data.FileTask.DEFAULT_VALUE_LENGTH_RANGE;

public class BatchDataTestUtils {

    private static Logger logger = LoggerFactory.getLogger(BatchDataTestUtils.class);

    private static Map<String, Class<? extends DataValidator>> dataTypeClasses = Map.of(
        "string", StringDataValidator.class,
        "set", SetDataValidator.class,
        "hash", HashDataValidator.class,
        "list", ListDataValidator.class,
        "zset", ZsetDataValidator.class,
        "stream", StreamDataValidator.class
    );

    public static DataValidatorRunner addValidators(final DataValidatorRunner tester, Map<String, File> typeFiles) {
        typeFiles.forEach((typeName, file) -> {
            String type = typeName.toLowerCase();
            Class<? extends DataValidator> validatorClass = dataTypeClasses.get(type);
            if (validatorClass != null) {
                tester.addValidator(validatorClass, file);
            } else {
                throw new IllegalArgumentException("Unsupported data type: " + type);
            }
        });

        return tester;
    }

    /**
     * 批量生成测试数据文件
     * @param typeFiles 数据类型和对应的输出文件
     * @param force 是否强制覆盖旧文件
     * @param dataCount 每种数据类型的key数量
     * @param subDataCount 每个key下的数据条数
     */
    public static void generateDataFiles(Map<String, File> typeFiles, boolean force,
                                         int dataCount, int subDataCount) {
        generateDataFiles(typeFiles, force, dataCount, subDataCount, null);
    }

    /**
     * 批量生成测试数据文件
     * @param typeFiles 数据类型和对应的输出文件
     * @param force 是否强制覆盖旧文件
     * @param dataCount 每种数据类型的key数量
     * @param subDataCount 每个key下的数据条数
     * @param valueLengthRange value长度范围，null表示使用默认值范围
     */
    public static void generateDataFiles(Map<String, File> typeFiles, boolean force,
                                           int dataCount, int subDataCount, int[] valueLengthRange) {

        TestDataGenerator dataGenerator = new TestDataGenerator();

        addDataFiles(dataGenerator, typeFiles, force, dataCount, subDataCount, valueLengthRange);


        int taskCount = dataGenerator.getTaskCount();
        if(taskCount > 0) {
            logger.info("Generating test data ...");
            int executed = dataGenerator.executeTasks();
            Assertions.assertEquals(taskCount, executed);
        }
        else {
            logger.info("Test data files already exists, skip generating ...");
        }

    }

    /**
     * 批量添加数据文件任务到数据生成器
     * @param dataGenerator 数据生成器实例
     * @param typeFiles 数据类型和对应的输出文件
     * @param force 是否强制覆盖旧文件
     * @param dataCount 每种数据类型的key数量
     * @param subDataCount 每个key下的数据条数
     */
    public static TestDataGenerator addDataFiles(final TestDataGenerator dataGenerator, Map<String, File> typeFiles, boolean force,
                                                 int dataCount, int subDataCount) {
        return addDataFiles(dataGenerator, typeFiles, force, dataCount, subDataCount, null);
    }

    /**
     * 批量添加数据文件任务到数据生成器
     * @param dataGenerator 数据生成器实例
     * @param typeFiles 数据类型和对应的输出文件
     * @param force 是否强制覆盖旧文件
     * @param dataCount 每种数据类型的key数量
     * @param subDataCount 每个key下的数据条数
     * @param valueLengthRange value长度范围，null表示使用默认值范围
     */
    public static TestDataGenerator addDataFiles(final TestDataGenerator dataGenerator, Map<String, File> typeFiles, boolean force,
                               int dataCount, int subDataCount, int[] valueLengthRange) {
        typeFiles.forEach((typeName, file) -> {
            String type = typeName.toLowerCase();
            Class<? extends DataValidator> validatorClass = dataTypeClasses.get(type);
            if (validatorClass != null) {
                if(force && file.exists()) {
                    file.delete(); // 强制覆盖旧文件
                }


                FileTask task = new FileTask(file, type, dataCount, subDataCount,
                        (null == valueLengthRange) ? DEFAULT_VALUE_LENGTH_RANGE : valueLengthRange);
                dataGenerator.addNoExistTask(task);
            } else {
                throw new IllegalArgumentException("Unsupported data type: " + type);
            }
        });

        return dataGenerator;
    }


    public static int delDataFiles(Map<String, File> typeFiles) {
        int deletedCount = 0;
        for (File file : typeFiles.values()) {
            if (file.exists() && file.delete()) {
                deletedCount++;
                logger.info("Deleted data file: {}", file.getAbsolutePath());
            }
        }

        return deletedCount;
    }



}
